import sys
import os
import cv2
from PIL import Image
import numpy as np
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QPushButton, QLabel, 
    QVBoxLayout, QHBoxLayout, QWidget, QLineEdit, 
    QFileDialog, QSlider, QMessageBox, QProgressBar
)
from PySide6.QtCore import Qt, QThread, Signal
from PySide6.QtGui import QFont, QPixmap, QImage

class VideoFileSelector(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("视频文件选择器")
        self.setGeometry(100, 100, 800, 500)
        
        # 视频相关变量
        self.video_path = ""
        self.cap = None
        self.total_frames = 0
        self.fps = 0
        
        # 创建中心部件和布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        
        # 文件选择区域
        file_layout = QHBoxLayout()
        self.select_button = QPushButton("选择视频文件")
        self.select_button.setMinimumHeight(40)
        self.select_button.setFont(QFont("SimHei", 10))
        self.select_button.clicked.connect(self.select_file)
        file_layout.addWidget(self.select_button)
        
        self.path_edit = QLineEdit()
        self.path_edit.setReadOnly(True)
        self.path_edit.setMinimumHeight(40)
        self.path_edit.setFont(QFont("SimHei", 10))
        file_layout.addWidget(self.path_edit, stretch=1)
        main_layout.addLayout(file_layout)
        
        # 视频信息标签
        self.info_label = QLabel("视频信息: 未加载")
        self.info_label.setFont(QFont("SimHei", 9))
        self.info_label.setAlignment(Qt.AlignLeft | Qt.AlignVCenter)
        main_layout.addWidget(self.info_label)
        
        # 操作按钮区域
        btn_layout = QHBoxLayout()
        
        self.load_button = QPushButton("加载视频")
        self.load_button.setMinimumHeight(40)
        self.load_button.setFont(QFont("SimHei", 10))
        self.load_button.clicked.connect(self.load_video)
        self.load_button.setEnabled(False)
        btn_layout.addWidget(self.load_button)
        
        self.save_frames_button = QPushButton("保存所有帧")
        self.save_frames_button.setMinimumHeight(40)
        self.save_frames_button.setFont(QFont("SimHei", 10))
        self.save_frames_button.clicked.connect(self.save_all_frames)
        self.save_frames_button.setEnabled(False)
        btn_layout.addWidget(self.save_frames_button)
        
        main_layout.addLayout(btn_layout)
        
        # 逐秒播放控制区域
        self.play_timer_layout = QHBoxLayout()
        
        self.play_button = QPushButton("逐秒播放")
        self.play_button.setMinimumHeight(40)
        self.play_button.setFont(QFont("SimHei", 10))
        self.play_button.clicked.connect(self.play_by_second)
        self.play_button.setEnabled(False)
        self.play_timer_layout.addWidget(self.play_button)
        
        self.time_slider = QSlider(Qt.Horizontal)
        self.time_slider.setRange(0, 0)
        self.time_slider.valueChanged.connect(self.update_frame_display)
        self.play_timer_layout.addWidget(self.time_slider)
        
        self.time_label = QLabel("时间: 00:00:00")
        self.time_label.setFont(QFont("SimHei", 10))
        self.time_label.setAlignment(Qt.AlignCenter)
        self.time_label.setMinimumWidth(120)
        self.play_timer_layout.addWidget(self.time_label)
        
        main_layout.addLayout(self.play_timer_layout)
        
        # 视频预览区域
        self.preview_label = QLabel("视频预览")
        self.preview_label.setAlignment(Qt.AlignCenter)
        self.preview_label.setMinimumSize(640, 360)
        self.preview_label.setStyleSheet("border: 1px solid #cccccc; background-color: #f0f0f0;")
        main_layout.addWidget(self.preview_label)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setTextVisible(True)
        self.progress_bar.hide()
        main_layout.addWidget(self.progress_bar)
        
        # 状态标签
        self.status_label = QLabel("就绪")
        self.status_label.setFont(QFont("SimHei", 10))
        self.status_label.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(self.status_label)
        
        # 设置布局间距
        main_layout.setSpacing(15)
        main_layout.setContentsMargins(20, 20, 20, 20)
        
        # 获取当前脚本所在路径
        self.script_dir = os.path.dirname(os.path.abspath(__file__))
        self.temp_dir = os.path.join(self.script_dir, "temp")
        
        # 创建保存帧的线程
        self.save_thread = None
    
    def select_file(self):
        """打开文件选择对话框并显示选择的文件路径"""
        file_filter = "视频文件 (*.mp4 *.avi *.mkv *.mov *.wmv);;所有文件 (*)"
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择视频文件", self.script_dir, file_filter
        )
        
        if file_path:
            self.video_path = file_path
            self.path_edit.setText(file_path)
            self.status_label.setText(f"已选择: {os.path.basename(file_path)}")
            self.load_button.setEnabled(True)
            self.save_frames_button.setEnabled(False)
            self.play_button.setEnabled(False)
            self.info_label.setText("视频信息: 未加载")
            self.preview_label.setText("视频预览")
        else:
            self.status_label.setText("未选择文件")
    
    def load_video(self):
        """加载视频文件并显示信息"""
        if not self.video_path:
            QMessageBox.warning(self, "警告", "请先选择视频文件")
            return
        
        self.cap = cv2.VideoCapture(self.video_path)
        if not self.cap.isOpened():
            QMessageBox.critical(self, "错误", "无法打开视频文件")
            return
        
        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.fps = self.cap.get(cv2.CAP_PROP_FPS)
        duration = self.total_frames / self.fps if self.fps > 0 else 0
        
        # 更新信息标签
        info = f"视频信息: 总帧数 {self.total_frames}, FPS {self.fps:.2f}, 时长 {duration:.2f}秒"
        self.info_label.setText(info)
        
        # 更新滑块范围
        self.time_slider.setRange(0, int(duration))
        self.time_slider.setValue(0)
        
        # 显示第一帧
        ret, frame = self.cap.read()
        if ret:
            self.display_frame(frame)
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)  # 重置到第一帧
        
        # 启用按钮
        self.save_frames_button.setEnabled(True)
        self.play_button.setEnabled(True)
        self.status_label.setText(f"已加载视频: {os.path.basename(self.video_path)}")
    
    def save_all_frames(self):
        """保存视频的所有帧到temp目录"""
        if not self.cap or not self.cap.isOpened():
            QMessageBox.warning(self, "警告", "请先加载视频")
            return
        
        # 检查并创建temp目录
        if not os.path.exists(self.temp_dir):
            os.makedirs(self.temp_dir)
            self.status_label.setText(f"创建临时目录: {self.temp_dir}")
        
        # 显示进度条
        self.progress_bar.show()
        self.progress_bar.setRange(0, self.total_frames)
        self.progress_bar.setValue(0)
        
        # 禁用按钮
        self.select_button.setEnabled(False)
        self.load_button.setEnabled(False)
        self.save_frames_button.setEnabled(False)
        self.play_button.setEnabled(False)
        
        # 创建并启动保存线程
        self.save_thread = SaveFramesThread(self.cap, self.temp_dir, self.total_frames)
        self.save_thread.progress_updated.connect(self.update_progress)
        self.save_thread.finished.connect(self.on_save_finished)
        self.save_thread.start()
    
    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)
    
    def on_save_finished(self, success, message):
        """保存完成后的回调"""
        # 恢复按钮状态
        self.select_button.setEnabled(True)
        self.load_button.setEnabled(True)
        self.save_frames_button.setEnabled(True)
        self.play_button.setEnabled(True)
        
        # 隐藏进度条
        self.progress_bar.hide()
        
        if success:
            self.status_label.setText(f"保存完成: {message}")
            QMessageBox.information(self, "成功", f"所有帧已保存到 {self.temp_dir}")
        else:
            self.status_label.setText(f"保存失败: {message}")
            QMessageBox.critical(self, "错误", f"保存帧时出错: {message}")
    
    def play_by_second(self):
        """逐秒播放视频"""
        if not self.cap or not self.cap.isOpened():
            QMessageBox.warning(self, "警告", "请先加载视频")
            return
        
        current_time = self.time_slider.value()
        if current_time >= int(self.total_frames / self.fps):
            current_time = 0  # 重置到开始
        
        # 计算对应的帧
        frame_pos = int(current_time * self.fps)
        self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_pos)
        
        # 读取并显示帧
        ret, frame = self.cap.read()
        if ret:
            self.display_frame(frame)
            self.time_slider.setValue(current_time + 1)  # 更新滑块
            self.update_time_label(current_time)
    
    def update_frame_display(self, value):
        """根据滑块值更新帧显示"""
        if not self.cap or not self.cap.isOpened():
            return
        
        # 计算对应的帧
        frame_pos = int(value * self.fps)
        self.cap.set(cv2.CAP_PROP_POS_FRAMES, frame_pos)
        
        # 读取并显示帧
        ret, frame = self.cap.read()
        if ret:
            self.display_frame(frame)
            self.update_time_label(value)
    
    def display_frame(self, frame):
        """显示帧到预览区域"""
        # 转换BGR为RGB
        rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # 调整图像大小以适应预览区域
        h, w, ch = rgb_image.shape
        bytes_per_line = ch * w
        q_img = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        
        # 缩放图像
        pixmap = QPixmap.fromImage(q_img)
        scaled_pixmap = pixmap.scaled(
            self.preview_label.size(), 
            Qt.KeepAspectRatio, 
            Qt.SmoothTransformation
        )
        
        # 显示图像
        self.preview_label.setPixmap(scaled_pixmap)
    
    def update_time_label(self, seconds):
        """更新时间标签显示"""
        hours = seconds // 3600
        minutes = (seconds % 3600) // 60
        secs = seconds % 60
        self.time_label.setText(f"时间: {hours:02d}:{minutes:02d}:{secs:02d}")

class SaveFramesThread(QThread):
    """保存视频帧的线程"""
    progress_updated = Signal(int)
    finished = Signal(bool, str)
    
    def __init__(self, cap, output_dir, total_frames):
        super().__init__()
        self.cap = cap
        self.output_dir = output_dir
        self.total_frames = total_frames
    
    def run(self):
        try:
            # 重置到第一帧
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
            
            for i in range(self.total_frames):
                ret, frame = self.cap.read()
                if not ret:
                    self.finished.emit(False, f"无法读取帧 {i}")
                    return
                
                # 使用PIL保存帧，解决中文路径问题
                if frame is not None:
                    # OpenCV使用BGR格式，PIL使用RGB格式
                    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    pil_img = Image.fromarray(rgb_frame)
                    
                    # 保存帧
                    frame_path = os.path.join(self.output_dir, f"frame_{i:06d}.jpg")
                    pil_img.save(frame_path)
                
                # 更新进度
                self.progress_updated.emit(i + 1)
            
            self.finished.emit(True, f"已保存 {self.total_frames} 帧")
        except Exception as e:
            self.finished.emit(False, str(e))

if __name__ == "__main__":
    app = QApplication(sys.argv)
    font = QFont("SimHei")
    app.setFont(font)
    window = VideoFileSelector()
    window.show()
    sys.exit(app.exec())